# encoding: utf-8

"""
模块描述

Authors: tongzhenguo
Date:    2023/01/19
"""
import numpy as np
import tensorflow as tf
from tensorflow.python import keras
from tensorflow.python.keras import layers
from keras.models import Sequential
from tensorflow.python.keras import callbacks


class MyModel(object):
    def __init__(self, input_shape):
        self.input_shape = input_shape
        self.model = Sequential(
            [
                layers.InputLayer(input_shape=self.input_shape),
                layers.SimpleRNN(units=128, return_sequences=True),
                layers.SimpleRNN(units=64, return_sequences=False),
                layers.Dense(1, activation='linear')
            ]
        )
        self.optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
        self.batch_size = 20
        self.epoches = 500

    def train(self, X, y):
        self.model.compile(optimizer=self.optimizer, loss="mse")
        callback_early_stopping = callbacks.EarlyStopping(
            monitor="loss",
            patience=10,
            min_delta=0.0001,
            restore_best_weights=True
        )
        his = self.model.fit(X, y,
                             batch_size=self.batch_size,
                             epochs=self.epoches,
                             callbacks=[callback_early_stopping])
        print(his)
        print(self.model.summary())
        return self.model


if __name__ == "__main__":
    from app.model.pandas_sample import Sample
    from collect_stock_data import StockBase

    base = StockBase()
    data = base.get_history_k_data(601318, '2022-01-01', '2023-01-10')
    sample = Sample()
    train_x, train_y, test_x, test_y = sample.generate(data)
    model = MyModel((train_x.shape[1], train_x.shape[2])).train(train_x, train_y)

    import matplotlib.pyplot as plt

    pred_y = model.predict(test_x)
    plt.figure(figsize=(15, 6))
    plt.plot(np.array(test_y))
    plt.plot(np.array(pred_y))
    plt.legend(['ground_truth', 'pred'])
    plt.show()

    # 打印原值
    print([sample.scaler.data_min_ + p * (sample.scaler.data_max_ - sample.scaler.data_min_) for p in pred_y])
